# 模板匹配（找图）
import cv2
import numpy as np
import os
import argparse
import pyautogui

def main(template_dir, left, top, width, height, threshold, unique_width, unique_height):
    # 指定要捕获的屏幕区域
    region = (left, top, width, height)

    # 1. 捕获屏幕
    original_image = pyautogui.screenshot(region = region)
    original_image = cv2.cvtColor(np.array(original_image), cv2.COLOR_RGB2BGR)
    
    # 2. 读取模板图片目录下的所有图片
    template_files = [f for f in os.listdir(template_dir) if f.endswith('.bmp')]

    # 3. 初始化计数器和识别结果
    match_count = 0
    results = ''

    # 4. 使用matchTemplate函数进行模板匹配
    for template_file in template_files:
        template_path = os.path.join(template_dir, template_file)
        template_image = cv2.imread(template_path)
        h, w = template_image.shape[:2]  # 获取模板的高度和宽度

        # 使用matchTemplate函数进行模板匹配
        res = cv2.matchTemplate(original_image, template_image, cv2.TM_CCOEFF_NORMED)

        # 找到所有超过阈值的匹配位置
        loc = np.where(res >= threshold)

        # 存储非重复的匹配位置
        unique_locs = []

        # 遍历所有匹配位置并画出矩形框
        for pt in zip(*loc[::-1]):  # *操作符用于解压元组，loc[::-1]是为了将(x, y)顺序调整为(y, x)
            # width, height = template_image.shape[1], template_image.shape[0]
            # print(f"Match found in '{template_file}' at: ({pt[0]}, {pt[1]}), Width: {width}, Height: {height}")

            # 检查这个点是否已经在unique_locs中存在
            is_unique = True
            for u in unique_locs:
                if abs(u[0] - pt[0]) <= unique_width and abs(u[1] - pt[1]) <= unique_height:
                    # 如果点太接近，则标记为重复
                    is_unique = False
                    break

            if is_unique:
                unique_locs.append(pt)
                match_count += 1  # 增加计数器
                result = f"{template_file},{pt[0]},{pt[1]}"
                results += f"|{result}" if results != f'' else result
                # cv2.rectangle(original_image, pt, (pt[0] + width, pt[1] + height), (0, 255, 0), 2)
    results = f"{match_count}:{results}"

    # 5. 输出结果和总的匹配数量
    print(results, flush=True)
    # print(f"Total number of unique matches found: {match_count}")

    # 6. 显示结果
    """ cv2.imshow('Match Result', original_image)
    cv2.waitKey(0)
    cv2.destroyAllWindows() """

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Capture and match templates from a specified screen region.")
    parser.add_argument("--template_dir", type=str, help="要识别的图片目录")
    parser.add_argument("--ltx", type=int, help="识别范围左上角横坐标", default=0)
    parser.add_argument("--lty", type=int, help="识别范围左上角纵坐标", default=0)
    parser.add_argument("--width", type=int, help="识别范围宽度", default=1920)
    parser.add_argument("--height", type=int, help="识别范围高度", default=1080)
    parser.add_argument("--threshold", type=float, help="相似度", default=0.8)
    parser.add_argument("--unique_width", type=int, help="横坐标差距小于这个值的结果会被去重", default=0)
    parser.add_argument("--unique_height", type=int, help="纵坐标差距小于这个值的结果会被去重", default=0)
    
    args = parser.parse_args()
    
    main(args.template_dir, args.ltx, args.lty, args.width, args.height, args.threshold, args.unique_width, args.unique_height)